# -*- coding: utf-8 -*-
# @Time    : 2024/9/7 22:24
# @Author  : Z.J. Zhang
# @Email   : zijingzhang@mail.ustc.edu.cn
# @File    : base.py
# @Software: PyCharm
import json
import os
import typing

import matplotlib

import myutils.filenametools

matplotlib.use('tkagg')
import matplotlib.pyplot as plt
import numpy
from scipy.interpolate import interpn
from abc import abstractmethod

def min_and_max(x):
    return numpy.array([min(x),max(x)])
class RectGrid3D:
    class Saver:
        key_all_pts = "all_points"
        key_V = "values"
        key_parameter_names = "parameter_names.json"

    def __init__(self, xs, ys, zs, parameter_names: typing.List[str]):
        self.xs, self.ys, self.zs = xs, ys, zs
        self.X, self.Y, self.Z = numpy.meshgrid(xs, ys, zs, indexing='ij')
        # self.parameter_names = parameter_names
        self.bounding_box = numpy.array([
            min_and_max(each) for each in [self.xs,self.ys,self.zs]
        ])
        self.parameter_names = []
        self.V = numpy.zeros((*self.X.shape, len(parameter_names)))

        self.add_new_parameter_names(parameter_names)
        # self.__dict_parameter_name_to_index = {pn: i for i, pn in enumerate(self.parameter_names )}
        self.get_index_of_parameter_name = numpy.vectorize(
            lambda parameter_name: self.__dict_parameter_name_to_index[parameter_name])

    def save(self, dir_path: str = None):
        if not dir_path: dir_path = "grid_%s" % (myutils.filenametools.get_unique_str())
        os.makedirs(dir_path, exist_ok=True)
        numpy.save(os.path.join(dir_path, self.Saver.key_all_pts), self.all_points())
        numpy.save(os.path.join(dir_path, self.Saver.key_V), self.V)
        with open(os.path.join(dir_path, self.Saver.key_parameter_names), "w") as f:
            json.dump(self.parameter_names, f)

    def all_points(self) -> numpy.ndarray:
        """

        :return: shape (..., 3)
        """
        return numpy.array((self.X, self.Y, self.Z)).transpose((1, 2, 3, 0))

    def get_parameter_data(self, parameter_names):
        """
        获取指定parameter_name的全部数据
        :param parameter_names:
        :return:
        """
        return self.V[:, :, :, self.get_index_of_parameter_name(parameter_names)]

    def get(self, x, y, z, parameter_names: typing.Iterable[str] = None):
        if parameter_names is None: parameter_names = self.parameter_names
        return interpn((self.xs, self.ys, self.zs), self.get_parameter_data(parameter_names), (x, y, z),bounds_error=False,fill_value=0.)

    def set_parameter_data(self, parameter_data: typing.Dict[str, numpy.ndarray]):
        """
        设置指定parameter_name的全部数据
        :param parameter_names:
        :return:
        """
        self.add_new_parameter_names(list(parameter_data.keys()))

        for parameter_name in parameter_data:
            self.V[:, :, :, self.get_index_of_parameter_name(parameter_name)] = parameter_data[parameter_name]

    def add_new_parameter_names(self, new_parameter_names: typing.List[str]):
        new_parameter_names = list(set(new_parameter_names) - set(self.parameter_names))
        if new_parameter_names:
            self.V = numpy.concatenate(
                (
                    self.V,
                    numpy.zeros((*self.X.shape, len(new_parameter_names)))
                ),
                axis=len(self.X.shape))
            self.parameter_names = self.parameter_names + new_parameter_names
        self.__dict_parameter_name_to_index = {pn: i for i, pn in enumerate(self.parameter_names)}


class FieldGenerator():
    @abstractmethod
    def cache(self, grid: RectGrid3D) -> None:  # RectGrid3D:
        return  # grid


default_grid = RectGrid3D(numpy.array((*numpy.linspace(-10e-3, -5e-3, 5),
                                       *numpy.linspace(-4.5e-3, 4.5e-3, 20),
                                       *numpy.linspace(5e-3, 10e-3, 5)
                                       )),
                          numpy.array((*numpy.linspace(-15e-3, -5e-3, 6),
                                       *numpy.linspace(-4.5e-3, 4.5e-3, 20),
                                       *numpy.linspace(5e-3, 10e-3, 5)
                                       )),
                          numpy.linspace(-10e-3, 10e-3, 22),
                          [
                              # "Bx", "By",                           "Bz"
                          ]
                          )

if __name__ == '__main__':
    plt.ion()


    def _test_func(x, y, z):
        r_ = 2e-3
        return numpy.exp(-(x ** 2 + y ** 2 + z ** 2) / (2 * r_ ** 2))


    grid = default_grid
    grid.set_parameter_data({grid.parameter_names[0]: _test_func(grid.X, grid.Y, grid.Z, )})
    grid.get(0, 0, 0)

    X, Y = numpy.meshgrid(grid.xs, grid.ys, indexing='ij')

    plt.figure()
    plt.contourf(X, Y, grid.get(X, Y, 0, ['Bx'])[:, :, 0])
    # plt.scatter(X,Y)
    plt.gca().set_aspect('equal')

    plt.figure()
    plt.plot(grid.xs, grid.get(grid.xs, 0, 0, ["Bx"])[:, 0])
